#!/usr/bin/env python3
"""
Enhanced WebSocket Handler - 增强的WebSocket消息处理器
专门处理EchoKit服务器的消息格式和机器人命令转换
"""

import asyncio
import json
import time
import threading
from typing import Dict, Any, Callable, Optional
import websockets
from websockets.exceptions import ConnectionClosed, InvalidURI, WebSocketException

from .echokit_message_handler import EchoKitMessageHandler


class EnhancedWebSocketHandler:
    """增强的WebSocket处理器，专门处理EchoKit服务器通信"""
    
    def __init__(self, url: str, logger=None, reconnect_interval: float = 5.0):
        self.url = url
        self.logger = logger
        self.reconnect_interval = reconnect_interval
        
        # WebSocket连接状态
        self.websocket = None
        self.connected = False
        self.running = False
        
        # 事件循环和线程管理
        self.loop = None
        self.thread = None
        
        # 消息处理器
        self.message_handler = EchoKitMessageHandler(logger)
        
        # 回调函数
        self.command_callback = None
        self.status_callback = None
        self.error_callback = None
        self.connection_callback = None
        
        # 统计信息
        self.message_count = 0
        self.last_message_time = None
        
    def set_command_callback(self, callback: Callable[[str, Dict], None]):
        """设置命令回调函数"""
        self.command_callback = callback
        
    def set_status_callback(self, callback: Callable[[str], None]):
        """设置状态回调函数"""
        self.status_callback = callback
        
    def set_error_callback(self, callback: Callable[[str], None]):
        """设置错误回调函数"""
        self.error_callback = callback
        
    def set_connection_callback(self, callback: Callable[[bool], None]):
        """设置连接状态回调函数"""
        self.connection_callback = callback
    
    def start(self):
        """启动WebSocket客户端"""
        if self.running:
            return
            
        self.running = True
        self.thread = threading.Thread(target=self._run_event_loop, daemon=True)
        self.thread.start()
        
        if self.logger:
            self.logger.info(f'Enhanced WebSocket handler started for {self.url}')
    
    def stop(self):
        """停止WebSocket客户端"""
        self.running = False
        
        if self.loop and self.loop.is_running():
            asyncio.run_coroutine_threadsafe(self._cleanup(), self.loop)
        
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=5.0)
            
        if self.logger:
            self.logger.info('Enhanced WebSocket handler stopped')
    
    def _run_event_loop(self):
        """运行异步事件循环"""
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        
        try:
            self.loop.run_until_complete(self._connect_with_retry())
        except Exception as e:
            if self.logger:
                self.logger.error(f'WebSocket event loop error: {e}')
        finally:
            self.loop.close()
    
    async def _connect_with_retry(self):
        """带重试的连接"""
        while self.running:
            try:
                await self._connect()
                if self.running:
                    if self.logger:
                        self.logger.info(f'Connection closed, retrying in {self.reconnect_interval}s')
                    await asyncio.sleep(self.reconnect_interval)
                    
            except (ConnectionClosed, WebSocketException, OSError) as e:
                if self.running:
                    if self.logger:
                        self.logger.warn(f'WebSocket connection error: {e}')
                    self._call_error_callback(f'Connection error: {e}')
                    await asyncio.sleep(self.reconnect_interval)
                    
            except Exception as e:
                if self.logger:
                    self.logger.error(f'Unexpected WebSocket error: {e}')
                if self.running:
                    await asyncio.sleep(self.reconnect_interval)
    
    async def _connect(self):
        """建立WebSocket连接"""
        try:
            if self.logger:
                self.logger.info(f'Connecting to {self.url}')
                
            async with websockets.connect(
                self.url,
                ping_interval=20,
                ping_timeout=10
            ) as websocket:
                self.websocket = websocket
                self.connected = True
                
                if self.logger:
                    self.logger.info('WebSocket connected successfully')
                self._call_connection_callback(True)
                self._call_status_callback('connected')
                
                # 消息处理循环
                async for message in websocket:
                    if not self.running:
                        break
                        
                    try:
                        await self._process_message(message)
                    except Exception as e:
                        if self.logger:
                            self.logger.error(f'Error processing message: {e}')
                            
        except InvalidURI as e:
            if self.logger:
                self.logger.error(f'Invalid WebSocket URI: {self.url} - {e}')
            self._call_error_callback(f'Invalid URI: {e}')
            
        except asyncio.TimeoutError:
            if self.logger:
                self.logger.warn('WebSocket connection timeout')
            self._call_error_callback('Connection timeout')
            
        finally:
            self.connected = False
            self.websocket = None
            self._call_connection_callback(False)
            if self.running:
                self._call_status_callback('disconnected')
    
    async def _process_message(self, raw_message):
        """处理接收到的消息"""
        try:
            self.message_count += 1
            self.last_message_time = time.time()
            
            if self.logger:
                self.logger.debug(f'Received raw message: {raw_message}')
            
            # 处理二进制消息（通常是音频数据）
            if isinstance(raw_message, bytes):
                if self.logger:
                    self.logger.debug(f'Received binary message of {len(raw_message)} bytes')
                return
            
            # 处理文本消息
            if isinstance(raw_message, str):
                # 使用消息处理器解析消息
                parsed_message = self.message_handler.parse_message(raw_message)
                
                if parsed_message:
                    await self._handle_parsed_message(parsed_message)
                else:
                    if self.logger:
                        self.logger.warn(f'Failed to parse message: {raw_message}')
                    self._call_error_callback(f'Failed to parse message: {raw_message}')
            
        except Exception as e:
            if self.logger:
                self.logger.error(f'Error in message processing: {e}')
            self._call_error_callback(f'Message processing error: {e}')
    
    async def _handle_parsed_message(self, parsed_message: Dict[str, Any]):
        """处理解析后的消息"""
        message_type = parsed_message.get('type', 'unknown')
        
        if message_type == 'voice_command':
            # 处理语音命令
            command = parsed_message.get('command')
            action = parsed_message.get('action')
            
            if command and action:
                if self.logger:
                    self.logger.info(f'Received voice command: {command} -> {action}')
                self._call_command_callback(command, action)
                self._call_status_callback(f'executing_{command}')
            else:
                if self.logger:
                    self.logger.warn(f'Invalid command: {command}')
                self._call_error_callback(f'Invalid command: {command}')
                
        elif message_type == 'text_message':
            # 处理文本消息
            content = parsed_message.get('content', '')
            if self.logger:
                self.logger.info(f'Received text message: {content}')
            self._call_status_callback(f'text_received')
            
        else:
            # 处理其他类型的消息
            if self.logger:
                self.logger.debug(f'Received message type: {message_type}')
            self._call_status_callback(f'message_{message_type}')
    
    def _call_command_callback(self, command: str, action: Dict[str, float]):
        """调用命令回调函数"""
        if self.command_callback:
            try:
                self.command_callback(command, action)
            except Exception as e:
                if self.logger:
                    self.logger.error(f'Error in command callback: {e}')
    
    def _call_status_callback(self, status: str):
        """调用状态回调函数"""
        if self.status_callback:
            try:
                self.status_callback(status)
            except Exception as e:
                if self.logger:
                    self.logger.error(f'Error in status callback: {e}')
    
    def _call_error_callback(self, error: str):
        """调用错误回调函数"""
        if self.error_callback:
            try:
                self.error_callback(error)
            except Exception as e:
                if self.logger:
                    self.logger.error(f'Error in error callback: {e}')
    
    def _call_connection_callback(self, connected: bool):
        """调用连接状态回调函数"""
        if self.connection_callback:
            try:
                self.connection_callback(connected)
            except Exception as e:
                if self.logger:
                    self.logger.error(f'Error in connection callback: {e}')
    
    async def send_message(self, message: Dict[str, Any]) -> bool:
        """发送消息到WebSocket服务器"""
        if not self.connected or not self.websocket:
            if self.logger:
                self.logger.warn('Cannot send message: WebSocket not connected')
            return False
            
        try:
            message_str = json.dumps(message)
            await self.websocket.send(message_str)
            
            if self.logger:
                self.logger.debug(f'Sent message: {message}')
            return True
            
        except Exception as e:
            if self.logger:
                self.logger.error(f'Failed to send message: {e}')
            return False
    
    def send_message_sync(self, message: Dict[str, Any]) -> bool:
        """同步发送消息"""
        if self.loop and self.loop.is_running():
            future = asyncio.run_coroutine_threadsafe(self.send_message(message), self.loop)
            try:
                return future.result(timeout=5.0)
            except asyncio.TimeoutError:
                if self.logger:
                    self.logger.error('Timeout sending message')
                return False
        return False
    
    async def _cleanup(self):
        """清理资源"""
        if self.websocket:
            try:
                await self.websocket.close()
            except Exception:
                pass
            self.websocket = None
        
        self.connected = False
    
    def is_connected(self) -> bool:
        """检查连接状态"""
        return self.connected
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        return {
            'connected': self.connected,
            'message_count': self.message_count,
            'last_message_time': self.last_message_time,
            'available_commands': self.message_handler.get_available_commands()
        }